"""
# -*- coding: utf-8 -*-
# @Time    : 2023/5/24 15:39
# @Author  : 王摇摆
# @FileName: Model_Manual.py
# @Software: PyCharm
# @Blog    ：https://blog.csdn.net/weixin_44943389?type=blog
"""

import numpy as np
from sklearn.tree import DecisionTreeRegressor


class gbdtmc:
    """
    梯度提升树多分类算法
    """

    def __init__(self, n_estimators=100, learning_rate=0.1):
        # 梯度提升树弱学习器数量
        self.n_estimators = n_estimators
        # 学习速率
        self.learning_rate = learning_rate
        print('人工GBDT多分类器已初始化完毕！')

    def fit(self, X, y):
        """
        梯度提升树多分类算法拟合
        """
        # 标签类，对应标签的数量
        self.y_classes, y_counts = np.unique(y, return_counts=True)
        # 标签类数量
        self.n_classes = len(self.y_classes)
        # 对标签进行One-Hot编码
        y_onehot = np.zeros((y.size, y.max() + 1))
        y_onehot[np.arange(y.size), y] = 1
        # 初始化 H0
        self.H0 = y_counts / X.shape[0]
        # 初始化预测值
        H = np.ones((X.shape[0], 1)).dot(self.H0.reshape(1, -1))
        # 估计器数组
        estimators = []
        # 叶子结点取值数组
        gammas = []
        # 遍历 n_estimators 次
        for k in range(self.n_estimators):
            H_exp = np.exp(H)
            H_exp_sum = np.sum(H_exp, axis=1)
            # 估计器
            sub_estimators = []
            # 叶子结点取值
            sub_gammas = []
            # 遍历 n_classes 次
            for m in range(self.n_classes):
                p_m = H_exp[:, m] / H_exp_sum
                # 计算第 m 个 y_hat
                y_hat_m = y_onehot[:, m] - p_m
                # 初始化决策回归树估计器
                estimator = DecisionTreeRegressor(max_depth=3, criterion="friedman_mse")
                # 将第 m 个 y_hat 当作标签值拟合训练集
                estimator.fit(X, y_hat_m)
                # 计算训练集在当前决策回归树的叶子结点
                leaf_ids = estimator.apply(X)
                # 每个叶子结点下包含的训练数据序号
                node_ids_dict = self.get_leaf_nodes(leaf_ids)
                # 叶子结点取值字典表
                gamma_dict = {}
                # 计算叶子结点取值
                for leaf_id, node_ids in node_ids_dict.items():
                    # 当前叶子结点包含的 y_hat
                    y_hat_sub = y_hat_m[node_ids]
                    y_hat_sub_abs = np.abs(y_hat_sub)
                    # 计算叶子结点取值
                    gamma = np.sum(y_hat_sub) / np.sum(np.multiply(y_hat_sub_abs, 1 - y_hat_sub_abs)) * (
                            self.n_classes - 1) / self.n_classes
                    gamma_dict[leaf_id] = gamma
                    # 更新预测值
                    H[node_ids, m] += self.learning_rate * gamma
                sub_estimators.append(estimator)
                sub_gammas.append(gamma_dict)
            estimators.append(sub_estimators)
            gammas.append(sub_gammas)
        self.estimators = estimators
        self.gammas = gammas

    def predict(self, X):
        """
        梯度提升树多分类算法预测
        """
        # 初始化预测值
        H = np.ones((X.shape[0], 1)).dot(self.H0.reshape(1, -1))
        # 遍历估计器
        for k in range(self.n_estimators):
            sub_estimators = self.estimators[k]
            sub_gammas = self.gammas[k]
            # 遍历分类数
            for m in range(self.n_classes):
                estimator = sub_estimators[m]
                # 计算在当前决策回归树的叶子结点
                leaf_ids = estimator.apply(X)
                # 每个叶子结点下包含的训练数据序号
                node_ids_dict = self.get_leaf_nodes(leaf_ids)
                # 叶子结点取值字典表
                gamma_dict = sub_gammas[m]
                # 计算预测值
                for leaf_id, node_ids in node_ids_dict.items():
                    gamma = gamma_dict[leaf_id]
                    H[node_ids, m] += self.learning_rate * gamma
        H_exp = np.exp(H)
        H_exp_sum = np.sum(H_exp, axis=1)
        # 计算概率
        probs = np.zeros((X.shape[0], self.n_classes))
        for m in range(self.n_classes):
            probs[:, m] = H_exp[:, m] / H_exp_sum
        return self.y_classes.take(np.argmax(probs, axis=1), axis=0)

    def get_leaf_nodes(self, leaf_ids):
        """
        每个叶子结点下包含的数据序号
        """
        node_ids_dict = {}
        for j in range(len(leaf_ids)):
            leaf_id = leaf_ids[j]
            node_ids = node_ids_dict.setdefault(leaf_id, [])
            node_ids.append(j)
        return node_ids_dict
